package ctp.statesgraph;

import java.util.Collection;
import java.util.HashSet;

import ctp.ctpgraph.CTPGraph;
import ctp.ctpgraph.Vertex;

public class PolicyBuilder {
	
	StatesGraph stGraph;
	Policy optimalPolicy = null;
	Policy optimalVerticesCommitingPolicy = null;

	public PolicyBuilder(StatesGraph stGraph) {
		super();
		this.stGraph = stGraph;
	}
	

	public Policy buildOptimalPolicy() {

		optimalPolicy = new Policy();
		optimalPolicy.setInitialState(stGraph.start);
		optimalPolicy(stGraph.start);
		return optimalPolicy;

		/*
		optimalPolicy = buildOptimalVerticesCommitingPolicy(new HashSet<Vertex>());
		return optimalPolicy;
		*/
	}
	
	private void optimalPolicy(State state) {
		if (optimalPolicy.getStateActionMap().containsKey(state)) {
			return;
		}
		
		if (state.isGoalState(stGraph.ctpGraph)) {
			PolicyAction pa = new PolicyAction(null, 0.0);
			optimalPolicy.getStateActionMap().put(state, pa);
		} else if (state instanceof ActionState) {
			Collection<Arc> possibleActions = this.stGraph.stGraph.getOutEdges(state);
			PolicyAction pa = new PolicyAction(null, Double.MAX_VALUE);			
			for (Arc arc : possibleActions) {
				State newState = this.stGraph.stGraph.getDest(arc);
				optimalPolicy(newState);
				double tempCost = ((Action)arc).cost + optimalPolicy.getStateActionMap().get(newState).cost; 
				if (tempCost < pa.cost) {
					pa.cost = tempCost;
					pa.action = ((Action)arc);
				}
			}
			optimalPolicy.getStateActionMap().put(state, pa);
		} else if (state instanceof ObservationState) {
			Collection<Arc> possibleObservations = this.stGraph.stGraph.getOutEdges(state);
			PolicyAction pa = new PolicyAction(null, 0);
			for (Arc arc : possibleObservations) {
				State newState = this.stGraph.stGraph.getDest(arc);
				optimalPolicy(newState);
				pa.cost = pa.cost + ((Observation)arc).getProb() * optimalPolicy.getStateActionMap().get(newState).cost;
			}
			optimalPolicy.getStateActionMap().put(state, pa);
		}
		
	}
	
	public Policy buildOptimalExploringVerticesPolicy(Collection<Vertex> explorationVertices) {
		Collection<Vertex> commitingVertices = new HashSet<Vertex>(stGraph.ctpGraph.getGraph().getVertices());
		commitingVertices.removeAll(explorationVertices);
		return buildOptimalVerticesCommitingPolicy(commitingVertices);
	}
	
	public Policy buildOptimalVerticesCommitingPolicy(Collection<Vertex> commitingVertices) {
		optimalVerticesCommitingPolicy = new Policy();
		optimalVerticesCommitingPolicy.setInitialState(stGraph.start);
		optimalVerticesCommitingPolicy(stGraph.start, commitingVertices);
		return optimalVerticesCommitingPolicy;
	}

	private void optimalVerticesCommitingPolicy(State state, Collection<Vertex> commitingVertices) {
		if (optimalVerticesCommitingPolicy.getStateActionMap().containsKey(state)) {
			return;
		}
		
		if (state.isGoalState(stGraph.ctpGraph)) {
			PolicyAction pa = new PolicyAction(null, 0.0);
			optimalVerticesCommitingPolicy.getStateActionMap().put(state, pa);
		} else if (state instanceof ActionState) {
			Collection<Arc> possibleActions = this.stGraph.stGraph.getOutEdges(state);
			PolicyAction pa = new PolicyAction(null, Double.MAX_VALUE);			
			for (Arc arc : possibleActions) {
				boolean actionIsCommiting = true;
				for (Vertex vertex : commitingVertices) {
					if (!Policy.isVertexStateActionCommiting(state, (Action)arc, vertex, stGraph.ctpGraph)) {
						actionIsCommiting = false;
						break;
					}
				}
				if (!actionIsCommiting) {
					continue;
				}
				State newState = this.stGraph.stGraph.getDest(arc);
				optimalVerticesCommitingPolicy(newState, commitingVertices);
				double tempCost = ((Action)arc).cost + optimalVerticesCommitingPolicy.getStateActionMap().get(newState).cost; 
				if (tempCost < pa.cost) {
					pa.cost = tempCost;
					pa.action = ((Action)arc);
				}
			}
			optimalVerticesCommitingPolicy.getStateActionMap().put(state, pa);
		} else if (state instanceof ObservationState) {
			Collection<Arc> possibleObservations = this.stGraph.stGraph.getOutEdges(state);
			PolicyAction pa = new PolicyAction(null, 0);
			for (Arc arc : possibleObservations) {
				State newState = this.stGraph.stGraph.getDest(arc);
				optimalVerticesCommitingPolicy(newState, commitingVertices);
				pa.cost = pa.cost + ((Observation)arc).getProb() * optimalVerticesCommitingPolicy.getStateActionMap().get(newState).cost;
			}
			optimalVerticesCommitingPolicy.getStateActionMap().put(state, pa);
		}
		
	}

	
}
